import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
from sklearn.metrics import calinski_harabasz_score

x, y = make_blobs(n_samples=500,
                  n_features=2,
                  centers=[[-1, -1], [0, 0], [1, 1], [2, 2]],
                  cluster_std=[0.4, 0.2, 0.2, 0.3],
                  random_state=22
                  )

plt.figure(figsize=(20, 10), dpi=80)
"""
K-Means
"""
y_pred = KMeans(n_clusters=4).fit_predict(x)  # KMeans聚类模型，参数4是将数据划分为4个簇

plt.scatter(x[:, 0], x[:, 1], c=y_pred)

plt.show()

print(calinski_harabasz_score(x,y_pred))
